from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detect.src.detect2file.predictor import VisualizationDemo
# constants
WINDOW_NAME = "COCO detections"
class DetectDemo():
    def __init__(self,config='detect/src/detect2file/faster_rcnn_R_50_C4_1x.yaml',weights='detect/src/detect2file/model_final_721ade.pkl'):
        self.config = config
        self.weights = weights
        cfg = get_cfg()
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
        cfg.merge_from_file(self.config)
        cfg.merge_from_list(["MODEL.WEIGHTS", self.weights])
        self.demo = VisualizationDemo(cfg)
    def detection(self,img):
        _, visualized_output = self.demo.run_on_image(img)
        return visualized_output
# a = DetectDemo()
# img = read_image('bochum_000000_002293_leftImg8bit_foggy_beta_0.005.png', format="BGR")
# out = a.detection(img)
# out.save('ss.jpg')